{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib.util\n",
    "\n",
    "if importlib.util.find_spec('cirq') is None:\n",
    "    print(\"installing cirq...\")\n",
    "    !pip install --quiet cirq\n",
    "    print(\"installed cirq.\")\n",
    "\n",
    "if importlib.util.find_spec('quimb') is None:\n",
    "    print(\"installing cirq-core[contrib]...\")\n",
    "    !pip install --quiet 'cirq-core[contrib]'\n",
    "    print(\"installed cirq-core[contrib].\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cirq to Tensor Networks\n",
    "\n",
    "Here we demonstrate turning circuits into tensor network representations of the circuit's unitary, final state vector, final density matrix, and final noisy density matrix. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cirq\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from cirq.contrib.svg import SVGCircuit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cirq.contrib.quimb as ccq\n",
    "import quimb.tensor as qtn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a random circuit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qubits = cirq.LineQubit.range(3)\n",
    "circuit = cirq.testing.random_circuit(qubits, n_moments=10, op_density=0.8, random_state=52)\n",
    "circuit = cirq.drop_empty_moments(circuit)\n",
    "SVGCircuit(circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Circuit to Tensors\n",
    "The circuit defines a tensor network representation. By default, the initial state is the `|0...0>` state (represented by the \"zero qubit\" operations labeled \"Q0\" in the legend. \"Q1\" are single qubit operations and \"Q2\" are two qubit operations. The open legs are the indices into the state vector and are of the form \"i{m}_q{n}\" where `m` is the time index (given by the returned `qubit_frontier` dictionary) and \"n\" is the qubit string."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tensors, qubit_frontier, fix = ccq.circuit_to_tensors(circuit, qubits)\n",
    "tn = qtn.TensorNetwork(tensors)\n",
    "print(qubit_frontier)\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "tn.graph(fix=fix, color=['Q0', 'Q1', 'Q2'], figsize=(8, 8))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### To dense"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi_tn = ccq.tensor_state_vector(circuit, qubits)\n",
    "psi_cirq = cirq.final_state_vector(circuit, qubit_order=qubits)\n",
    "np.testing.assert_allclose(psi_cirq, psi_tn, atol=1e-7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Circuit Unitary\n",
    "We can also leave the input legs open which gives a tensor network representation of the unitary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tensors, qubit_frontier, fix = ccq.circuit_to_tensors(circuit, qubits, initial_state=None)\n",
    "tn = qtn.TensorNetwork(tensors)\n",
    "print(qubit_frontier)\n",
    "tn.graph(fix=fix, color=['Q0', 'Q1', 'Q2'], figsize=(8, 8))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### To dense"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_tn = ccq.tensor_unitary(circuit, qubits)\n",
    "u_cirq = circuit.unitary(qubit_order=qubits)\n",
    "np.testing.assert_allclose(u_cirq, u_tn, atol=1e-7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Density Matrix\n",
    "We can also turn a circuit into its density matrix. The density matrix resulting from the evolution of the `|0><0|` initial state can be thought of as two copies of the circuit: one going \"forwards\" and one going \"backwards\" (i.e. use the complex conjugate of each operation). Kraus operator noise operations \"link\" the forwards and backwards circuits. As such, the density matrix for pure states is simple.\n",
    "\n",
    "Note: for density matrices, we return a `fix` variable for a circuit-like layout of the tensors when calling `tn.graph`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tensors, qubit_frontier, fix = ccq.circuit_to_density_matrix_tensors(circuit=circuit, qubits=qubits)\n",
    "tn = qtn.TensorNetwork(tensors)\n",
    "tn.graph(fix=fix, color=['Q0', 'Q1', 'Q2'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Noise\n",
    "Noise operations entangle the forwards and backwards evolutions. The new tensors labeled \"kQ1\" are 1-qubit Kraus operators."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "noise_model = cirq.ConstantQubitNoiseModel(cirq.DepolarizingChannel(p=1e-3))\n",
    "circuit = cirq.Circuit(noise_model.noisy_moments(circuit.moments, qubits))\n",
    "SVGCircuit(circuit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tensors, qubit_frontier, fix = ccq.circuit_to_density_matrix_tensors(circuit=circuit, qubits=qubits)\n",
    "tn = qtn.TensorNetwork(tensors)\n",
    "tn.graph(fix=fix, color=['Q0', 'Q1', 'Q2', 'kQ1'], figsize=(8, 8))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### For 6 or fewer qubits, we specify the contraction ordering.\n",
    "For low-qubit-number circuits, a reasonable contraction ordering is to go in moment order (as a normal simulator would do). Otherwise, quimb will try to find an optimal ordering which was observed to take longer than it takes to do the contraction itself. We show how to tell quimb to contract in order by using the moment tags."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "partial = 12\n",
    "tags_seq = [(f'i{i}b', f'i{i}f') for i in range(partial)]\n",
    "tn.graph(fix=fix, color=[x for x, _ in tags_seq] + [y for _, y in tags_seq], figsize=(8, 8))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The result of a partial contraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tn2 = tn.contract_cumulative(tags_seq, inplace=False)\n",
    "tn2.graph(fix=fix, color=['Q0', 'Q1', 'Q2', 'kQ1'], figsize=(8, 8))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### To Dense"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rho_tn = ccq.tensor_density_matrix(circuit, qubits)\n",
    "rho_cirq = cirq.final_density_matrix(circuit, qubit_order=qubits)\n",
    "np.testing.assert_allclose(rho_cirq, rho_tn, atol=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Profile\n",
    "For low-qubit-number, deep, noisy circuits, the quimb contraction is faster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import timeit\n",
    "\n",
    "\n",
    "def profile(n_qubits: int, n_moments: int):\n",
    "    qubits = cirq.LineQubit.range(n_qubits)\n",
    "    circuit = cirq.testing.random_circuit(qubits, n_moments=n_moments, op_density=0.8)\n",
    "    noise_model = cirq.ConstantQubitNoiseModel(cirq.DepolarizingChannel(p=1e-3))\n",
    "    circuit = cirq.Circuit(noise_model.noisy_moments(circuit.moments, qubits))\n",
    "    circuit = cirq.drop_empty_moments(circuit)\n",
    "    n_moments = len(circuit)\n",
    "    variables = {'circuit': circuit, 'qubits': qubits}\n",
    "\n",
    "    setup1 = ['import cirq', 'import numpy as np']\n",
    "    n_call_cs, duration_cs = timeit.Timer(\n",
    "        stmt='cirq.final_density_matrix(circuit)', setup='; '.join(setup1), globals=variables\n",
    "    ).autorange()\n",
    "\n",
    "    setup2 = ['from cirq.contrib.quimb import tensor_density_matrix', 'import numpy as np']\n",
    "    n_call_t, duration_t = timeit.Timer(\n",
    "        stmt='tensor_density_matrix(circuit, qubits)', setup='; '.join(setup2), globals=variables\n",
    "    ).autorange()\n",
    "\n",
    "    return {\n",
    "        'n_qubits': n_qubits,\n",
    "        'n_moments': n_moments,\n",
    "        'duration_cirq': duration_cs,\n",
    "        'duration_quimb': duration_t,\n",
    "        'n_call_cirq': n_call_cs,\n",
    "        'n_call_quimb': n_call_t,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "records = []\n",
    "max_qubits = 6\n",
    "max_moments = 500\n",
    "for n_qubits in [3, max_qubits]:\n",
    "    for n_moments in range(1, max_moments, 50):\n",
    "        record = profile(n_qubits=n_qubits, n_moments=n_moments)\n",
    "        records.append(record)\n",
    "        print('.', end='', flush=True)\n",
    "\n",
    "df = pd.DataFrame(records)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def select(df, k, v):\n",
    "    return df[df[k] == v].drop(k, axis=1)\n",
    "\n",
    "\n",
    "pd.DataFrame.select = select\n",
    "\n",
    "\n",
    "def plot1(df, labelfmt):\n",
    "    for k in ['duration_cirq', 'duration_quimb']:\n",
    "        plt.plot(df['n_moments'], df[k], '.-', label=labelfmt.format(k))\n",
    "    plt.legend(loc='best')\n",
    "\n",
    "\n",
    "def plot(df):\n",
    "    df['duration_cirq'] /= df['n_call_cirq']\n",
    "    df['duration_quimb'] /= df['n_call_quimb']\n",
    "    plot1(df.select('n_qubits', 3), 'n = 3, {}')\n",
    "    plot1(df.select('n_qubits', 6), 'n = 6, {}')\n",
    "    plt.xlabel('N Moments')\n",
    "    plt.ylabel('Time / s')\n",
    "\n",
    "\n",
    "plot(df)\n",
    "plt.tight_layout()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
